import paddle
import paddle.nn as nn

class FallDetectionLSTM(nn.Layer):
    def __init__(self, input_size=4, hidden_size=32, num_layers=1, num_classes=2):
        super(FallDetectionLSTM, self).__init__()
        # LSTM 层：处理时序数据 (Batch, Time_Steps, Features)
        self.lstm = nn.LSTM(input_size=input_size, 
                            hidden_size=hidden_size, 
                            num_layers=num_layers, 
                            direction='forward')
        # 全连接层：分类 (跌倒 vs 正常)
        self.fc = nn.Linear(hidden_size, num_classes)
        self.softmax = nn.Softmax()

    def forward(self, x):
        # x shape: [batch_size, time_steps, input_size]
        # LSTM 输出: (output, (h_n, c_n))
        output, (h_n, c_n) = self.lstm(x)
        # 取最后一个时间步的输出用于分类
        last_hidden = h_n[-1] 
        logits = self.fc(last_hidden)
        return self.softmax(logits)

if __name__ == "__main__":
    # 测试网络结构
    model = FallDetectionLSTM()
    paddle.summary(model, (1, 30, 4)) # 假设 30帧数据，每帧4个特征
    print("模型结构定义测试通过！")